#!/usr/bin/env python
#-*- coding: utf-8 -*-

# Author:   Edward Roualdes
#           University of Kentucky

from __future__ import print_function
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from collections import defaultdict
import utils

# function definition
def sortBarcode(INPUTFILE, BARCODEFILE, trim):
    """Sort DNA reads from INPUTFILE by individual's names in BARCODEFILE searching first for the barcode in the beginning 30 nucleotides and then, if previously unsuccesful, the reverse complement's first 30 nucleotides.

    input: 
    sorted            = list of tuples [(barcode_1, dnaRead_1), ... (barcode_n, dnaRead_n)]
    n                 = number of DNA reads
    nomatch           = number of reads with no matchin barcodes
    multmatch         = number of reads that matched multiple barcodes
    match             = number of reads that were matched

    output:
    c                 = number of reads sorted by barcode
    sortedBarcodes    = defaultdict( (barcode_name_1, dnaRead_1_1, ..., dnaRead_1_m), ...
    (barcode_name_c, dnaRead_c_1, ..., dnaRead_c_l) )
"""


    sortedBarcodes = defaultdict(list) # initialize the data structure to be returned

    # create list of barcode strings from BARCODEFILE
    # list parsing less memory efficient when we want all info per sequence
    bOpen = open(BARCODEFILE, "rU")
    bList = list(SeqIO.parse(bOpen, "fasta"))
    bOpen.close()
    # print(len(bList))

    # create list of dna reads from INPUTFILE
    iOpen = open(INPUTFILE, "rU")
    iList = list(SeqIO.parse(iOpen, "fasta"))
    iOpen.close()
    # print(len(iList))

    # initializations
    n = len(iList)              # number of reads
    nomatch = 0
    multmatch = 0

    # match reads to barcodes
    for read in iList:          # for each read
        match = 0
        sequence = read.seq # temporary copy read
        l = len(sequence)   # length of read
        if l > 30:
            s = sequence.lower()[0:29] # first 30 nucleotides in read
            r = sequence.reverse_complement().lower()[0:29] # first 30 rev_comp
            for code in bList:
                bs = code.seq.lower() # barcode
                lb = len(bs)          # length barcode
                found = s.find(bs) # find index of bs in seq if present ow -1
                if found > -1:     # if barcode found within first 30
                    if trim:
                        # trim beg/end by found+lb
                        # read.seq = sequence[found+lb : l-found-lb]
                        # trim beg/end by 30
                        read.seq = sequence[30:l-30] # trim 30
                    else:
                        read.seq = sequence
                    match += 1
                    whichBarcode = code.name
                    break
                else:
                    found = r.find(bs) # find index of bs in rev_seq
                    if found > -1:
                        if trim:
                            # trim beg/end by found+lb
                            # read.seq = sequence[found+lb:l-found-lb]
                            # trim beg/end by 30
                            read.seq = sequence[30:l-30] # trim 30
                        else:
                            read.seq = sequence
                        match += 1
                        whichBarcode = code.name
                        break
            if match == 1:
                sortedBarcodes[whichBarcode].append(read)
            elif match == 0:
                nomatch += 1
            # else:
            #     multmatch += 1
        else:
            nomatch += 1


    c = n - nomatch - multmatch # no. BARCODE classified reads

    del iList, bList
    print("Classified", c, "reads by barcode...")
    return sortedBarcodes, c


# if reads are already sorted by individual and contained in INPUTFILE
# this function intelligently skips over sortBarcode
# thus eliminating the need for a BARCODEFILE
def sortInds(indsFile):
    u"""from sortedInds create a structure identical to that which sortBarcode creates.

Input:
    indsFile = file containing reads enumerated by individual's name underscore read_number: edward_j 

Output: (re sortBarcode)
    c                 = number of reads sorted by barcode
    sortedBarcodes    = defaultdict( (barcode_1, dnaRead_1_1, ..., dnaRead_1_m), ...
    (barcode_c, dnaRead_c_1, ..., dnaRead_c_l) )
    """

    # import reads
    with open(indsFile, "rU") as f:
        reads = list(SeqIO.parse(f, "fasta"))    

    # get individuals' names
    names = list(set( r.name[:r.name.rfind('_')] for r in reads ))

    # put reads into defaultdict
    # arranged by individuals' names
    c = 0
    sortedReads = defaultdict(list)
    for r in reads:
        for n in names:
            if n in r.name:
                sortedReads[n].append(r)
                c += 1          # number of reads
                break           # move on to next read if successfully sorted

    del reads
    print("Imported", c, "reads...")
    return sortedReads, c

# if __name__ == "__main__":
